package com.szm.sentinelx.starter.prometheus;

import com.alibaba.csp.sentinel.metric.extension.MetricExtension;
import com.alibaba.csp.sentinel.slots.block.BlockException;

import io.prometheus.client.CollectorRegistry;
import io.prometheus.client.Counter;
import io.prometheus.client.Gauge;
import io.prometheus.client.Histogram;


/**
 * @author: shaozeming
 * @date: 2022/9/14 20:09
 * @description:
 **/
public class PrometheusMetricLoader implements MetricExtension {

    private Counter passRequests;
    private Counter blockRequests;
    private Counter successRequests;
    private Counter exceptionRequests;
    private Histogram rtHist;
    private Gauge currentThreads;

    private final CollectorRegistry collectorRegistry;


    public PrometheusMetricLoader(CollectorRegistry collectorRegistry){
        this.collectorRegistry= collectorRegistry;
        init();
    }

    private void init(){
        passRequests = Counter.build()
                .name("sentinel_pass_requests_total")
                .help("total pass requests.")
                .labelNames("resource")
                .register(collectorRegistry);
        blockRequests = Counter.build()
                .name("sentinel_block_requests_total")
                .help("total block requests.")
                .labelNames("resource", "type", "ruleLimitApp", "limitApp")
                .register(collectorRegistry);
        successRequests = Counter.build()
                .name("sentinel_success_requests_total")
                .help("total success requests.")
                .labelNames("resource")
                .register(collectorRegistry);
        exceptionRequests = Counter.build()
                .name("sentinel_exception_requests_total")
                .help("total exception requests.")
                .labelNames("resource")
                .register(collectorRegistry);
        currentThreads = Gauge.build()
                .name("sentinel_current_threads")
                .help("current thread count.")
                .labelNames("resource")
                .register(collectorRegistry);
        rtHist = Histogram.build()
                .name("sentinel_requests_latency_seconds")
                .help("request latency in seconds.")
                .labelNames("resource")
                .register(collectorRegistry);
        PrometheusMetric.setMetricExtension(this);
    }

    /**
     * Add current pass count of the resource name.
     *
     * @param resource resource name
     * @param n        count to add
     * @param args     additional arguments of the resource, eg. if the resource is a method name,
     *                 the args will be the parameters of the method.
     */
    @Override
    public void addPass(String resource, int n, Object... args) {
       passRequests.labels(resource).inc(n);
    }

    /**
     * Add current block count of the resource name.
     *
     * @param resource       resource name
     * @param n              count to add
     * @param origin         the original invoker.
     * @param blockException block exception related.
     * @param args           additional arguments of the resource, eg. if the resource is a method name,
     *                       the args will be the parameters of the method.
     */
    @Override
    public void addBlock(String resource, int n, String origin, BlockException blockException, Object... args) {
        blockRequests.labels(resource).inc(n);
    }

    /**
     * Add current completed count of the resource name.
     *
     * @param resource resource name
     * @param n        count to add
     * @param args     additional arguments of the resource, eg. if the resource is a method name,
     *                 the args will be the parameters of the method.
     */
    @Override
    public void addSuccess(String resource, int n, Object... args) {
        successRequests.labels(resource).inc(n);
    }

    /**
     * Add current exception count of the resource name.
     *
     * @param resource  resource name
     * @param n         count to add
     * @param throwable exception related.
     */
    @Override
    public void addException(String resource, int n, Throwable throwable) {
        exceptionRequests.labels(resource).inc(n);
    }

    /**
     * Add response time of the resource name.
     *
     * @param resource resource name
     * @param rt       response time in millisecond
     * @param args     additional arguments of the resource, eg. if the resource is a method name,
     *                 the args will be the parameters of the method.
     */
    @Override
    public void addRt(String resource, long rt, Object... args) {
        rtHist.labels(resource).observe((double) rt);
    }

    /**
     * Increase current thread count of the resource name.
     *
     * @param resource resource name
     * @param args     additional arguments of the resource, eg. if the resource is a method name,
     *                 the args will be the parameters of the method.
     */
    @Override
    public void increaseThreadNum(String resource, Object... args) {
       currentThreads.labels(resource).inc();
    }

    /**
     * Decrease current thread count of the resource name.
     *
     * @param resource resource name
     * @param args     additional arguments of the resource, eg. if the resource is a method name,
     *                 the args will be the parameters of the method.
     */
    @Override
    public void decreaseThreadNum(String resource, Object... args) {
        currentThreads.labels(resource).dec();
    }

}
